
import asyncio
import websockets

from lj_tool import tool_async, tool_id
from . import tool_log


class LJWs():
    def __init__(self, port, on_msg, **options) -> None:
        self.port = port
        self.on_msg = on_msg
        self.conns = {}
        self.__dict__.update(options)

    def start(self):
        async def init_start():
            async with websockets.serve(
                    lambda conn, path: self.handle_msg(conn, path), "0.0.0.0", self.port, ping_interval=None):
                tool_log.info(f"LJ - websocket_server[{self.port}]启动成功...")
                await asyncio.Future()
                tool_log.info(f"LJ - websocket_server[{self.port}]关闭成功...")
        tool_async.async_start_run(init_start())

    def get_conn(self, conn_id):
        return self.conns.get(conn_id)

    async def send(self, conn_id, msg):
        conn = self.conns.get(conn_id)
        print(f"LJ - ws发送消息[{msg[0:30]}]...")
        if conn:
            try:
                await conn.send(msg)
            except:
                self.close(conn_id)

    def close(self, conn_id):
        tool_log.info(f"LJ - 连接[{conn_id}]断开...")
        # 这里增加判定，有时候没有处理
        if self.conns.get(conn_id):
            del self.conns[conn_id]
        if self.on_close:
            self.on_close(conn_id)

    async def handle_msg(self, conn, path):
        conn_id = tool_id.gen_id()
        tool_log.info(f"LJ - 新连接[{path}]进入...")

        self.conns[conn_id] = conn

        if self.on_conn:
            # await tool_async.async_thread_run_2(lambda: self.on_conn(conn_id, path))
            await self.on_conn(conn_id, path)

        try:
            async for msg in conn:
                self.on_msg(conn_id, msg)
        finally:
            self.close(conn_id)
